import bettermoments as bm
import os
import numpy as np
from astropy.io import fits
path='V883_Ori_SB_HDO-225.535GHz_robust_2.0.image.fits'
data, velax = bm.load_cube(path)
rms = bm.estimate_RMS(data=data, N=5)
user_mask = bm.get_user_mask(data=data, user_mask_path='V883_Ori_SB_HDO-225.535GHz_robust_2.0.mask-hdo.fits')
channel_mask=bm.get_channel_mask(data=data,firstchannel=28,lastchannel=49)
threshold_mask = bm.get_threshold_mask(data=data,clip=-100.0,smooth_threshold_mask=0.0)
mask = bm.get_combined_mask(threshold_mask=threshold_mask,
                            channel_mask=channel_mask,
                            user_mask=user_mask,
                            combine='and')
masked_data = data * mask
moments = bm.collapse_zeroth(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='zeroth', path=path)
moments = bm.collapse_first(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='first', path=path)
moments = bm.collapse_second(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='second', path=path)
moments_mask = bm.collapse_zeroth(velax=velax, data=mask, rms=0.0)
bm.save_to_FITS(moments=moments_mask, method='zeroth', path='V883_Ori_SB_HDO-225.535GHz_robust_2.0.mask-hdo.fits')



#HDO 241
path='V883_Ori_SB_HDO-241.558GHz_robust_2.0.image.fits'
data, velax = bm.load_cube(path)
rms = bm.estimate_RMS(data=data, N=5)
user_mask = bm.get_user_mask(data=data, user_mask_path='V883_Ori_SB_HDO-241.558GHz_robust_2.0.mask-hdo.fits')
channel_mask=bm.get_channel_mask(data=data,firstchannel=31,lastchannel=49)
threshold_mask = bm.get_threshold_mask(data=data,clip=-100.0,smooth_threshold_mask=0.0)
mask = bm.get_combined_mask(threshold_mask=threshold_mask,
                            channel_mask=channel_mask,
                            user_mask=user_mask,
                            combine='and')
masked_data = data * mask
moments = bm.collapse_zeroth(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='zeroth', path=path)
moments = bm.collapse_first(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='first', path=path)
moments = bm.collapse_second(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='second', path=path)
moments_mask = bm.collapse_zeroth(velax=velax, data=mask, rms=0.0)
bm.save_to_FITS(moments=moments_mask, method='zeroth', path='V883_Ori_SB_HDO-241.558GHz_robust_2.0.mask-hdo.fits')



#H2CO
path='V883_Ori_SB_H2CO-225.694GHz-fullcube_robust_2.0.image.fits'
data, velax = bm.load_cube(path)
smoothed_data = bm.smooth_data(data=data, smooth=3, polyorder=0)
rms = bm.estimate_RMS(data=smoothed_data, N=5)
#user_mask = bm.get_user_mask(data=data, #user_mask_path='V883_Ori_SB_HDO-241.558GHz_robust_2.0.image.fits')
channel_mask=bm.get_channel_mask(data=smoothed_data,firstchannel=212,lastchannel=260)
threshold_mask = bm.get_threshold_mask(data=smoothed_data,clip=-5.0,smooth_threshold_mask=0.0)
#user_mask = bm.get_user_mask(data=smoothed_data, user_mask_path=None)
user_mask = bm.get_user_mask(data=smoothed_data, user_mask_path='V883_Ori_SB_H2CO-225.694GHz-fullcube_robust_2.0.mask-h2co.fits')
mask = bm.get_combined_mask(threshold_mask=threshold_mask,
                            channel_mask=channel_mask,
                            user_mask=user_mask,
                            combine='and')

masked_data = smoothed_data * mask
moments = bm.collapse_zeroth(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]*-1.0/1000.0,moments[1]*-1.0/1000.0)
bm.save_to_FITS(moments=moments, method='zeroth', path=path)
moments = bm.collapse_first(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='first', path=path)
moments = bm.collapse_second(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='second', path=path)
moments_mask = np.abs(bm.collapse_zeroth(velax=velax, data=np.abs(mask), rms=0.0))
bm.save_to_FITS(moments=moments_mask, method='zeroth', path='V883_Ori_SB_H2CO-225.694GHz-fullcube_robust_2.0.mask-h2co.fits')

#C17O
path='V883_Ori_SB_C17O-224.711GHz_robust_2.0.image.fits'
data, velax = bm.load_cube(path)
rms = bm.estimate_RMS(data=data, N=5)
#user_mask = bm.get_user_mask(data=data, #user_mask_path='V883_Ori_SB_HDO-241.558GHz_robust_2.0.image.fits')
channel_mask=bm.get_channel_mask(data=data,firstchannel=72,lastchannel=120)
threshold_mask = bm.get_threshold_mask(data=data,clip=-5.0,smooth_threshold_mask=0.0)
#user_mask = bm.get_user_mask(data=data, user_mask_path=None)
user_mask = bm.get_user_mask(data=data, user_mask_path='V883_Ori_SB_C17O-224.711GHz_robust_2.0.mask-c17o.fits')
mask = bm.get_combined_mask(threshold_mask=threshold_mask,
                            channel_mask=channel_mask,
                            user_mask=user_mask,
                            combine='and')

masked_data = data * mask
moments = bm.collapse_zeroth(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='zeroth', path=path)
moments = bm.collapse_first(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='first', path=path)
moments = bm.collapse_second(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='second', path=path)
moments_mask =np.abs( bm.collapse_zeroth(velax=velax, data=np.abs(mask), rms=0.0))
bm.save_to_FITS(moments=moments_mask, method='zeroth', path='V883_Ori_SB_C17O-224.711GHz_robust_2.0.mask-c17o.fits')

#CH3OH
lines=['241.806','241.813','241.852','241.879','241.887']
chan_ranges=[[96,147],[155,204],[474,531],[691,744],[761,814]]
path='V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.image.fits'
data, velax = bm.load_cube(path)
rms = bm.estimate_RMS(data=data, N=5)
for i in range(len(lines)):
   #user_mask = bm.get_user_mask(data=data, #user_mask_path='V883_Ori_SB_HDO-241.558GHz_robust_2.0.image.fits')
   channel_mask=bm.get_channel_mask(data=data,firstchannel=chan_ranges[i][0],lastchannel=chan_ranges[i][1])
   threshold_mask = bm.get_threshold_mask(data=data,clip=-5.0,smooth_threshold_mask=0.0)
   #user_mask = bm.get_user_mask(data=data, user_mask_path=None)
   user_mask = bm.get_user_mask(data=data, user_mask_path='V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.mask-ch3oh-'+lines[i]+'.fits')
   mask = bm.get_combined_mask(threshold_mask=threshold_mask,
                               channel_mask=channel_mask,
                               user_mask=user_mask,
                               combine='and')

   masked_data = data * mask
   moments = bm.collapse_zeroth(velax=velax, data=masked_data, rms=rms)
   moments=(moments[0]*-1.0/1000.0,moments[1]*-1.0/1000.0)
   bm.save_to_FITS(moments=moments, method='zeroth', path=path)
   moments = bm.collapse_first(velax=velax, data=masked_data, rms=rms)
   moments=(moments[0]/1000.0,moments[1]/1000.0)
   bm.save_to_FITS(moments=moments, method='first', path=path)
   moments = bm.collapse_second(velax=velax, data=masked_data, rms=rms)
   moments=(moments[0]/1000.0,moments[1]/1000.0)
   bm.save_to_FITS(moments=moments, method='second', path=path)
   for moment in ['M0','M1','M2','dM0','dM1','dM2']:
      os.system('mv V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.image_'+moment+'.fits V883_Ori_SB_CH3OH-241.846GHz-fullcube_robust_2.0.image_'+lines[i]+'_'+moment+'.fits')

#H218O
path='V883_Ori_SB_H218O-203.GHz-0.4kms_robust_2.0.image.fits'
data, velax = bm.load_cube(path)
rms = bm.estimate_RMS(data=data, N=5)
rms=0.00158
user_mask = bm.get_user_mask(data=data, user_mask_path='V883_Ori_SB_H218O-203.GHz-0.4kms_robust_2.0.mask-h218o.fits')
channel_mask=bm.get_channel_mask(data=data,firstchannel=31,lastchannel=49)
threshold_mask = bm.get_threshold_mask(data=data,clip=-100.0,smooth_threshold_mask=0.0)
mask = bm.get_combined_mask(threshold_mask=threshold_mask,
                            channel_mask=channel_mask,
                            user_mask=user_mask,
                            combine='and')
mask=channel_mask*user_mask
hdul = fits.open(path)
hdul[0].data=mask
hdul.writeto('h218o-mask-M0.fits')

masked_data = data * mask
moments = bm.collapse_zeroth(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='zeroth', path=path)
moments_mask = bm.collapse_zeroth(velax=velax, data=mask, rms=0.0)
bm.save_to_FITS(moments=moments_mask, method='zeroth', path='V883_Ori_SB_H218O-203.GHz-0.4kms_robust_2.0.mask-h218o.fits')

moments = bm.collapse_first(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='first', path=path)
moments = bm.collapse_second(velax=velax, data=masked_data, rms=rms)
moments=(moments[0]/1000.0,moments[1]/1000.0)
bm.save_to_FITS(moments=moments, method='second', path=path)


